import pandas as pd
import numpy as np
from chat_bot.models.intent_fasttext.utils import extact_field, get_input
from chat_bot.models.intent_fasttext import config
from sklearn.utils import shuffle
from sklearn.model_selection import StratifiedKFold


# 处理为可以输入的形式
def dataProcess(data_dir, new_data_dir):
    # 读取数据集
    data = pd.read_csv(data_dir)

    # 数据预处理
    # 提取label、query
    label = extact_field(data, 'label_class')
    query = extact_field(data, 'text')

    # 处理label格式
    label_list = []
    for list in label:
        new_list = "__label__" + list
        label_list.append((new_list))

    # 实体替代和停用词
    query_list = []
    for text in query:
        res = get_input(text)
        query_list.append(res)

    # 构建10折交叉验证训练和测试数据集
    sfolder = StratifiedKFold(n_splits=10)
    label = label_list
    text = query_list
    for index, (train_index, test_index) in enumerate(sfolder.split(label, text)):
        train_label, train_text = np.array(query_list)[train_index], np.array(label_list)[train_index]
        test_label, test_text = np.array(query_list)[test_index], np.array(label_list)[test_index]

        train_data = pd.DataFrame(zip(train_text, train_label))
        train_data = shuffle(train_data)
        train_data.to_csv(f'./data/train{index}.csv', index=False, header=['label', 'text'], sep='\t')
        test_data = pd.DataFrame(zip(test_text, test_label))
        test_data = shuffle(test_data)
        test_data.to_csv(f'./data/test{index}.csv', index=False, header=['label', 'text'], sep='\t')

    # 合成需要全部新数据并放入csv文件中
    new_data = []
    for i in range(0, len(query_list)):
        new_data.append(label_list[i] + " " + query_list[i])
    new_data = shuffle(new_data)
    new = pd.DataFrame(data=new_data)
    new.to_csv(new_data_dir, encoding='utf-8', index=False, header=False, sep='\t')


if __name__ == '__main__':
    dataProcess(config.merge_data_path, config.all_data_path)
    print("完成数据处理")
